package org.aurora.framework.boot.mybatis.flex.processor;

import com.mybatisflex.annotation.Column;
import com.mybatisflex.annotation.ColumnAlias;
import com.mybatisflex.annotation.Table;
import org.apache.velocity.Template;
import org.apache.velocity.VelocityContext;
import org.apache.velocity.app.VelocityEngine;
import org.apache.velocity.runtime.resource.loader.ClasspathResourceLoader;
import org.aurora.framework.boot.mybatis.flex.processor.config.GenerateConfig;
import org.aurora.framework.boot.mybatis.flex.processor.data.ClassInfo;
import org.aurora.framework.boot.mybatis.flex.processor.data.ColumnInfo;
import org.aurora.framework.boot.mybatis.flex.processor.utils.FileUtil;
import org.aurora.framework.boot.mybatis.flex.processor.utils.StringUtil;

import javax.annotation.processing.AbstractProcessor;
import javax.annotation.processing.ProcessingEnvironment;
import javax.annotation.processing.RoundEnvironment;
import javax.lang.model.SourceVersion;
import javax.lang.model.element.*;
import javax.lang.model.type.DeclaredType;
import javax.lang.model.type.TypeKind;
import javax.lang.model.type.TypeMirror;
import javax.lang.model.util.Elements;
import javax.lang.model.util.Types;
import java.io.File;
import java.io.IOException;
import java.io.Writer;
import java.math.BigDecimal;
import java.math.BigInteger;
import java.sql.Time;
import java.sql.Timestamp;
import java.time.*;
import java.time.chrono.JapaneseDate;
import java.time.format.DateTimeFormatter;
import java.util.*;

/**
 * name: CodeGenerationProcessor
 * <p>
 * desc: 查询处理器
 * </p>
 * create date: 2024-10-08 16:44:46
 *
 * @author cao.yong
 */
public class CodeGenerationProcessor extends AbstractProcessor {

    private static final List<String> DEFAULT_SUPPORT_COLUMN_TYPES = Arrays.asList(
        int.class.getName(), Integer.class.getName(),
        short.class.getName(), Short.class.getName(),
        long.class.getName(), Long.class.getName(),
        float.class.getName(), Float.class.getName(),
        double.class.getName(), Double.class.getName(),
        boolean.class.getName(), Boolean.class.getName(),
        Date.class.getName(), java.sql.Date.class.getName(), Time.class.getName(), Timestamp.class.getName(),
        Instant.class.getName(), LocalDate.class.getName(), LocalDateTime.class.getName(), LocalTime.class.getName(),
        OffsetDateTime.class.getName(), OffsetTime.class.getName(), ZonedDateTime.class.getName(),
        Year.class.getName(), Month.class.getName(), YearMonth.class.getName(), JapaneseDate.class.getName(),
        byte[].class.getName(), Byte[].class.getName(), Byte.class.getName(),
        BigInteger.class.getName(), BigDecimal.class.getName(),
        char.class.getName(), String.class.getName(), Character.class.getName()
    );

    private Types typeUtils;
    private Elements elementUtils;
    private GenerateConfig config;

    @Override
    public synchronized void init(ProcessingEnvironment processingEnv) {
        super.init(processingEnv);
        this.config = GenerateConfig.load(processingEnv.getFiler());
        this.typeUtils = processingEnv.getTypeUtils();
        this.elementUtils = processingEnv.getElementUtils();
    }

    @Override
    public boolean process(Set<? extends TypeElement> annotations, RoundEnvironment roundEnv) {
        if (!config.isEnabled()) {
            return true;
        }
        // 初始化模板引擎
        VelocityEngine engine = new VelocityEngine();
        engine.setProperty(VelocityEngine.RESOURCE_LOADER, "classpath");
        engine.setProperty("classpath.resource.loader.class", ClasspathResourceLoader.class.getName());
        engine.init();
        for (TypeElement annotation : annotations) {
            if (annotation.getQualifiedName().toString().equals(Table.class.getName())) {
                for (Element element : roundEnv.getElementsAnnotatedWith(annotation)) {
                    Table tableAnnotation = element.getAnnotation(Table.class);
                    String tableName = tableAnnotation.value();
                    String tableComment = tableAnnotation.comment();
                    String entityPackage = element.getEnclosingElement().toString();
                    String entityClassName = element.getSimpleName().toString();
                    String entityName = entityClassName.replace(config.getEntity().getSuffix(), "");
                    String tableDefEntityName = StringUtil.convertFieldName(entityName, config.getNamingStrategy());
                    String tableDefPackageName = entityPackage.substring(0, entityPackage.lastIndexOf(".")) + ".table";
                    String tableDefinePackage = config.getEntity().getTableDefinedPackage();
                    if (!StringUtil.isBlank(tableDefinePackage)) {
                        tableDefPackageName = tableDefinePackage;
                    }
                    String tableDefClassName = entityName + config.getEntity().getTableDefinedSuffix();
                    try {
                        String now = DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss").format(LocalDateTime.now());
                        ClassInfo classInfo = new ClassInfo(tableName, tableDefEntityName, tableDefPackageName, tableDefClassName, tableComment, config.getAuthor(), now);
                        // 循环获取直到无父类
                        TypeElement entityElement = (TypeElement) element;
                        do {
                            // 获取类属性和默认查询字段
                            fillColumnList(classInfo, entityElement, tableAnnotation.camelToUnderline());
                            entityElement = (TypeElement) typeUtils.asElement(entityElement.getSuperclass());
                        } while (entityElement != null);
                        // 数据表定义模板生成
                        VelocityContext tableDefContext = new VelocityContext();
                        tableDefContext.put("classInfo", classInfo);
                        tableDefContext.put("namingStrategy", config.getNamingStrategy());
                        tableDefContext.put("StrTool", StringUtil.class);

                        Writer writer = processingEnv.getFiler().createSourceFile(tableDefPackageName + "." + tableDefClassName).openWriter();
                        Template tpl = engine.getTemplate("templates/EntityQuery.vm", "UTF-8");
                        tpl.merge(tableDefContext, writer);
                        writer.close();

                        // 数据访问Mapper模板生成
                        if (config.getMapper().isEnabled()) {
                            String mapperPackage = entityPackage.substring(0, entityPackage.lastIndexOf("."))  + ".mapper";
                            String basePackage = config.getMapper().getBasePackage();
                            if (!StringUtil.isBlank(basePackage)) {
                                mapperPackage = basePackage;
                            }
                            String mapperClassName = entityName + config.getMapper().getSuffix();
                            String classFileName = mapperClassName + ".java";
                            VelocityContext mapperContext = new VelocityContext();
                            mapperContext.put("packageName", mapperPackage);
                            mapperContext.put("className", mapperClassName);
                            mapperContext.put("classComment", tableComment);
                            mapperContext.put("entityPackage", entityPackage);
                            mapperContext.put("entityName", entityClassName);
                            mapperContext.put("author", config.getAuthor());
                            mapperContext.put("createDate", now);
                            Writer fileWriter = FileUtil.createFileWriter(processingEnv.getFiler(), mapperPackage.replace(".", File.separator), classFileName);
                            if (fileWriter != null) {
                                Template tplJava = engine.getTemplate("templates/EntityMapper.vm", "UTF-8");
                                tplJava.merge(mapperContext, fileWriter);
                                fileWriter.close();
                            }
                        }
                    } catch (IOException e) {
                        throw new RuntimeException(e);
                    }
                }
            }
        }
        return false;
    }

    @Override
    public Set<String> getSupportedAnnotationTypes() {
        Set<String> annotations = new LinkedHashSet<>();
        annotations.add(Table.class.getCanonicalName());
        return annotations;
    }

    @Override
    public SourceVersion getSupportedSourceVersion() {
        return SourceVersion.latestSupported();
    }

    /**
     * 填充字段信息
     * @param classInfo 类信息
     * @param entityElement 基础元素
     * @param camelToUnderline 是否将驼峰命名转换为下划线命名
     */
    private void fillColumnList(ClassInfo classInfo, TypeElement entityElement, boolean camelToUnderline) {
        // 遍历所有包含的字段元素
        for (Element fieldElement : entityElement.getEnclosedElements()) {
            if (ElementKind.FIELD == fieldElement.getKind()) {
                Set<Modifier> modifiers = fieldElement.getModifiers();
                if (modifiers.contains(Modifier.STATIC)) {
                    // ignore static fields
                    continue;
                }
                Column column = fieldElement.getAnnotation(Column.class);
                if (column != null && column.ignore()) {
                    continue;
                }
                // 获取 typeHandlerClass 的名称，通过 column.typeHandler() 获取会抛出异常：MirroredTypeException:
                // 参考 https://stackoverflow.com/questions/7687829/java-6-annotation-processing-getting-a-class-from-an-annotation
                final String[] typeHandlerClass = {""};
                List<? extends AnnotationMirror> mirrors = fieldElement.getAnnotationMirrors();
                for (AnnotationMirror mirror : mirrors) {
                    mirror.getElementValues().forEach((executableElement, annotationValue) -> {
                        if ("typeHandler".contentEquals(executableElement.getSimpleName())) {
                            typeHandlerClass[0] = annotationValue.toString();
                        }
                    });
                }
                TypeMirror typeMirror = fieldElement.asType();
                Element element = this.typeUtils.asElement(typeMirror);
                if (element != null) {
                    typeMirror = element.asType();
                }
                String typeString = typeMirror.toString().trim();

                TypeElement typeElement = null;
                if (typeMirror.getKind() == TypeKind.DECLARED) {
                    typeElement = (TypeElement) ((DeclaredType) typeMirror).asElement();
                }

                // 未配置 typeHandler 的情况下，只支持基本数据类型，不支持比如 list set 或者自定义的类等
                if ((column == null || "org.apache.ibatis.type.UnknownTypeHandler".equals(typeHandlerClass[0]))
                    && !DEFAULT_SUPPORT_COLUMN_TYPES.contains(typeString)
                    && (typeElement != null && ElementKind.ENUM != typeElement.getKind())
                ) {
                    continue;
                }

                String property = fieldElement.toString();
                String columnName;
                if (column != null && !StringUtil.isBlank(column.value())) {
                    columnName = column.value();
                } else {
                    if (camelToUnderline) {
                        columnName = StringUtil.camelToUnderline(property);
                    } else {
                        columnName = property;
                    }
                }

                String[] aliases = getColumnAliasesByGetterMethod(entityElement, property);
                if (aliases == null || aliases.length == 0) {
                    ColumnAlias columnAlias = fieldElement.getAnnotation(ColumnAlias.class);
                    if (columnAlias != null) {
                        aliases = columnAlias.value();
                    }
                }
                ColumnInfo columnInfo = new ColumnInfo();
                columnInfo.setProperty(property);
                columnInfo.setColumn(columnName);
                columnInfo.setAliases(aliases);
                columnInfo.setComment(elementUtils.getDocComment(fieldElement));

                classInfo.addColumn(columnInfo);

                if (column == null || (!column.isLarge() && !column.isLogicDelete())) {
                    classInfo.addQueryColumn(columnName);
                }
            }
        }
    }

    private String[] getColumnAliasesByGetterMethod(TypeElement entityElement, String property) {
        if (null == entityElement) {
            return null;
        }
        // 遍历所有包含的方法元素
        for (Element method : entityElement.getEnclosedElements()) {
            if (ElementKind.METHOD == method.getKind()) {
                String methodName = method.toString();
                if (StringUtil.isGetterMethod(methodName, property)) {
                    ColumnAlias alias = method.getAnnotation(ColumnAlias.class);
                    if (alias != null) {
                        return alias.value();
                    } else {
                        // 重写方法，忽略别名
                        return null;
                    }
                }
            }
        }
        return getColumnAliasesByGetterMethod((TypeElement) this.typeUtils.asElement(entityElement.getSuperclass()), property);
    }
}
